from matplotlib import pyplot
import torch
import numpy


def Analysis_weights():
    device = torch.device('cuda')
    weights = torch.load('mnist0.lenet5.torch.weights')

    keys = tuple(weights)
    for key in keys:
        print(key, weights[key].shape)
    ana_w = weights['conv1.weight'].view(6, 5, 5)
    n_w = numpy.array(ana_w)


    fig = pyplot.figure()
    for i in range(6):
        data = n_w[i]
        ax = fig.add_subplot(331 + i)
        im = ax.imshow(data,cmap=pyplot.get_cmap('bwr'))
        pyplot.colorbar(im)
    if True:
        data = n_w[i]
        data[0][0] = -1.0
        data[0][1] = 1.0
        ax = fig.add_subplot(331 + 6)
        im = ax.imshow(data, cmap=pyplot.get_cmap('bwr'))
        pyplot.colorbar(im)


    pyplot.show()


def ana2():
    weights = torch.load('mnist0.lenet5.torch.weights')
    keys = tuple(weights)
    for k in keys:

        ana_w = weights[k]
        n_w = numpy.array(ana_w)
        n_f=n_w.flatten()
        print(n_f.shape)
        pyplot.hist(n_f)
        pyplot.show()


if __name__ == '__main__':
    ana2()
